/******************************************************************************
 * Project: Processor Affinity Manager                                        *
 * File:    Manager.cs                                                        *
 * Author:  Aaron Jones                                                       *
 * Version: 1                                                                 *
 *                                                                            *
 * Description:                                                               *
 *    This file contains the implementation for the main management gui for   *
 *    the processor affinity management.  The gui assists in managing the     *
 *    processor affinities of existing processes, and also helps maintain     *
 *    settings for default processor affinities.                              *
 ******************************************************************************/

using IWshRuntimeLibrary;
using System;
using System.Diagnostics;
using System.Windows.Forms;

namespace ProcessorAffinityManager
{
    public partial class Manager : Form
    {
        private const string InaccessibleProcess = "---";
        private ConfigurationFile m_ConfigFile = new ConfigurationFile();

        public Manager()
        {
            InitializeComponent();
            CurrentProcessSetAffinityButton.Enabled = false;
            RefreshCurrentProcessList();

            ProcessDefaultsEditButton.Enabled = false;
            ProcessDefaultsRemoveButton.Enabled = false;

            m_ConfigFile.Changed += new ConfigurationFile.ConfigurationChangedEvent(ConfigFileChanged);
            ConfigFileChanged();

            if (ShellExtension.Exists)
                ShellExtensionCheckBox.Checked = true;
            if (Service.Running)
                WindowsServiceCheckBox.Checked = true;

        }

        delegate void UpdateDefaultListView();
        void ConfigFileChanged()
        {
            if (ProcessDefaultsListView.InvokeRequired)
                Invoke(new UpdateDefaultListView(ConfigFileChanged));
            else
            {
                ProcessDefaultsListView.Items.Clear();
                foreach (Default d in m_ConfigFile.Defaults)
                {
                    ListViewItem item = new ListViewItem(new string[] { d.ExecutablePath, AffinityToString((IntPtr)d.ProcessorAffinity) });
                    item.Tag = d;

                    ProcessDefaultsListView.Items.Add(item);
                }
            }
        }

        private void SetCurrentProcessAffinityClick(object sender, EventArgs e)
        {
            Process proc = (Process)(CurrentProcessesListView.SelectedItems[0].Tag);
            Picker affinityPicker = new Picker((int)proc.ProcessorAffinity);
            if (affinityPicker.ShowDialog() == DialogResult.OK)
            {
                try
                {
                    proc.ProcessorAffinity = (IntPtr)affinityPicker.ProcessorAffinity;
                }

                catch
                {
                    MessageBox.Show("Access denied!  You do not have the proper privilages to change the affinity of this task.");
                }
            }
            RefreshCurrentProcessList();
        }

        private void RefreshCurrentProcessListClick(object sender, EventArgs e)
        {
            RefreshCurrentProcessList();
        }

        private void RefreshCurrentProcessList()
        {
            CurrentProcessesListView.Items.Clear();
            foreach (Process p in Process.GetProcesses())
            {
                ListViewItem item;
                IntPtr affinity = (IntPtr)0;

                try
                {
                    affinity = p.ProcessorAffinity;
                }

                catch
                {
                }

                item = new ListViewItem(
                    new string[] { p.ProcessName, AffinityToString(affinity) });
                item.Tag = p; // Easy access to change the affinity later.
                CurrentProcessesListView.Items.Add(item);
            }
        }

        private string AffinityToString(IntPtr affinity)
        {
            if (affinity == (IntPtr)0)
                return "---";
            else
            {
                string affinityString = "";
                for (int i = 1; i <= Environment.ProcessorCount; i++)
                    if (((int)affinity & (1 << i - 1)) != 0)
                        affinityString += i.ToString() + ",";
                return affinityString.Remove(affinityString.Length - 1);
            }
        }

        private void SelectCurrentProcess(object sender, EventArgs e)
        {
            CurrentProcessSetAffinityButton.Enabled = false;

            if (CurrentProcessesListView.SelectedIndices.Count == 1)
            {
                if(CurrentProcessesListView.SelectedItems[0].SubItems[1].Text != InaccessibleProcess)
                    CurrentProcessSetAffinityButton.Enabled = true;
            }
        }

        private void AddDefaultClick(object sender, EventArgs e)
        {
            Picker affinityPicker = new Picker();
            if (ProcessDefaultsOpenDialog.ShowDialog() == DialogResult.OK)
                if (affinityPicker.ShowDialog() == DialogResult.OK)
                    m_ConfigFile.AddDefault(
                        new Default(ProcessDefaultsOpenDialog.FileName, affinityPicker.ProcessorAffinity));

            ProcessDefaultsListView.SelectedItems.Clear();
        }

        private void EditDefaultClick(object sender, EventArgs e)
        {
            Default def = (Default)(ProcessDefaultsListView.SelectedItems[0].Tag);
            Picker affinityPicker = new Picker(def.ProcessorAffinity);
            if (affinityPicker.ShowDialog() == DialogResult.OK)
            {
                def.ProcessorAffinity = (int)affinityPicker.ProcessorAffinity;
                m_ConfigFile.AddDefault(def);
            }

            ProcessDefaultsListView.SelectedItems.Clear();
        }

        private void RemoveDefaultClick(object sender, EventArgs e)
        {
            Default def = (Default)(ProcessDefaultsListView.SelectedItems[0].Tag);
            m_ConfigFile.RemoveDefault(def);
            ProcessDefaultsListView.SelectedItems.Clear();
        }

        private void SelectProcessDefault(object sender, EventArgs e)
        {
            if (ProcessDefaultsListView.SelectedIndices.Count == 1)
            {
                ProcessDefaultsEditButton.Enabled = true;
                ProcessDefaultsRemoveButton.Enabled = true;
            }

            else
            {
                ProcessDefaultsEditButton.Enabled = false;
                ProcessDefaultsRemoveButton.Enabled = false;
            }
        }

        private void ShellExtensionChecked(object sender, EventArgs e)
        {
            if (ShellExtensionCheckBox.Checked)
                ShellExtension.Install();
            else
                ShellExtension.Remove();
        }

        private void ShortcutSaveClick(object sender, EventArgs e)
        {
            Picker affinityPicker = new Picker();
            if (affinityPicker.ShowDialog() == DialogResult.OK)
            {
                if (ShortcutTargetOpenDialog.ShowDialog() == DialogResult.OK)
                {
                    if (ShortcutSaveDialog.ShowDialog() == DialogResult.OK)
                    {
                        try
                        {
                            WshShell shell = new WshShell();
                            IWshShortcut link = (IWshShortcut)shell.CreateShortcut(ShortcutSaveDialog.FileName + ".lnk");
                            link.TargetPath = ConfigurationFile.InstallDirectory + Shortcut.ExecutableName;
                            link.Arguments = ((int)(affinityPicker.ProcessorAffinity)).ToString() + " " + ShortcutTargetOpenDialog.FileName;
                            link.Save();
                        }

                        catch
                        {
                            MessageBox.Show("Could not create the shortcut!");
                        }
                    }
                }
            }
        }

        private void DragNDropSaveClick(object sender, EventArgs e)
        {
            Picker affinityPicker = new Picker();
            if (affinityPicker.ShowDialog() == DialogResult.OK)
            {
                if (DragNDropSaveDialog.ShowDialog() == DialogResult.OK)
                {
                    try
                    {
                        WshShell shell = new WshShell();
                        IWshShortcut link = (IWshShortcut)shell.CreateShortcut(DragNDropSaveDialog.FileName + ".lnk");
                        link.TargetPath = ConfigurationFile.InstallDirectory + Shortcut.ExecutableName;
                        link.Arguments = ((int)(affinityPicker.ProcessorAffinity)).ToString();
                        link.Save();
                    }

                    catch
                    {
                        MessageBox.Show("Could not create the shortcut!");
                    }
                }
            }
        }

        private void WindowsServiceChecked(object sender, EventArgs e)
        {
            try
            {
                if (WindowsServiceCheckBox.Checked)
                    Service.Enable();
                else
                    Service.Disable();
            }

            catch
            {
                MessageBox.Show("There was an error changing the service status!");
                WindowsServiceCheckBox.Checked = !WindowsServiceCheckBox.Checked;
            }
        }
    }

    static class ManagerEntry
    {
        /// <summary>
        /// The main entry point for the application.
        /// </summary>
        [STAThread]
        static void Main()
        {
            try
            {
                Application.EnableVisualStyles();
                Application.SetCompatibleTextRenderingDefault(false);
                Application.Run(new Manager());
            }

            catch
            {
                MessageBox.Show("Fatal Error, the application will now close.");
            }
        }
    }
}
